#网络结构查看./ASimpleCNN.png
import torch
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

#一个手写的CNN
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        #定义两个卷积层
        self.conv1 = torch.nn.Conv2d(1,10,kernel_size=5)#第一个卷积层 输入通道是1 输出通道是10
        self.conv2 = torch.nn.Conv2d(10,20,kernel_size=6)#第一个卷积层 输入通道是10 输出通道是20
        self.pooling = torch.nn.MaxPool2d(2)#定义核为2*2的池化层
        self.fc = torch.nn.Linear(320,10)#定义全连接层 输入通道为320 输出通道为10

    def forward(self,x):
        #Flatten data from (n,1,28,28) tp (n,784)
        batch_size = x.size(0)#batchsize为n  x.size(0)返回值为参数为size的索引为0的参数，即n
        x = F.relu(self.pooling(self.con1(x)))#卷积—>池化->relu
        x = F.relu(self.pooling(self.con2(x)))#卷积—>池化->relu
        x = x.view(batch_size,-1)#flatten
        x = self.fc(x)
        return x

model = Net()

#启用gpu
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)#把模型转入到cuda里




